// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <algorithm>
#include <cstddef>
#include <memory>
#include <string>
#include <tuple>
#include <type_traits>
#include <utility>
#include <vector>

#include <nebula/core/array.h>
#include <nebula/array/builder_base.h>
#include <nebula/array/builder_binary.h>
#include <nebula/array/builder_nested.h>
#include <nebula/array/builder_primitive.h>
#include <nebula/core/chunked_array.h>
#include <nebula/compute/api.h>
#include <turbo/utility/status.h>
#include <nebula/core/table.h>
#include <nebula/types/type_fwd.h>
#include <nebula/types/type_traits.h>
#include <turbo/base/checked_cast.h>
#include <turbo/base/macros.h>

namespace nebula {

    class Schema;

    namespace stl {

        namespace internal {

            template<typename T, typename = void>
            struct is_optional_like : public std::false_type {
            };

            template<typename T, typename = void>
            struct is_dereferencable : public std::false_type {
            };

            template<typename T>
            struct is_dereferencable<T, nebula::internal::void_t<decltype(*std::declval<T>())>>
                    : public std::true_type {
            };

            template<typename T>
            struct is_optional_like<
                    T, typename std::enable_if<
                            std::is_constructible<bool, T>::value && is_dereferencable<T>::value &&
                            !std::is_array<typename std::remove_reference<T>::type>::value>::type>
                    : public std::true_type {
            };

            template<size_t N, typename Tuple>
            using BareTupleElement =
                    typename std::decay<typename std::tuple_element<N, Tuple>::type>::type;

        }  // namespace internal

        template<typename T, typename R = void>
        using enable_if_optional_like =
                typename std::enable_if<internal::is_optional_like<T>::value, R>::type;

        /// Traits meta class to map standard C/C++ types to equivalent Nebula types.
        template<typename T, typename Enable = void>
        struct ConversionTraits {
        };

        /// Returns builder type for given standard C/C++ type.
        template<typename CType>
        using CBuilderType =
                typename TypeTraits<typename ConversionTraits<CType>::ArrowType>::BuilderType;

        /// Default implementation of append_list_values.
        ///
        /// This function can be specialized by user to take advantage of appending
        /// contiguous ranges while appending. This default implementation will call
        /// ConversionTraits<ValueCType>::append_row() for each value in the range.
        template<typename ValueCType, typename Range>
        inline turbo::Status append_list_values(CBuilderType<ValueCType> &value_builder,
                                              Range &&cell_range) {
            for (auto const &value: cell_range) {
                TURBO_RETURN_NOT_OK(ConversionTraits<ValueCType>::append_row(value_builder, value));
            }
            return turbo::OkStatus();
        }

#define NEBULA_STL_CONVERSION(CType_, ArrowType_)                                    \
      template <>                                                                       \
      struct ConversionTraits<CType_> : public CTypeTraits<CType_> {                    \
        static turbo::Status append_row(typename TypeTraits<ArrowType_>::BuilderType& builder,  \
                                CType_ cell) {                                          \
          return builder.append(cell);                                                  \
        }                                                                               \
        static CType_ get_entry(const typename TypeTraits<ArrowType_>::ArrayType& array, \
                               size_t j) {                                              \
          return array.value(j);                                                        \
        }                                                                               \
      };                                                                                \
                                                                                        \
      template <>                                                                       \
      inline turbo::Status append_list_values<CType_, const std::vector<CType_>&>(               \
          typename TypeTraits<ArrowType_>::BuilderType & value_builder,                 \
          const std::vector<CType_>& cell_range) {                                      \
        return value_builder.append_values(cell_range);                                  \
      }

        NEBULA_STL_CONVERSION(bool, BooleanType)

        NEBULA_STL_CONVERSION(int8_t, Int8Type)

        NEBULA_STL_CONVERSION(int16_t, Int16Type)

        NEBULA_STL_CONVERSION(int32_t, Int32Type)

        NEBULA_STL_CONVERSION(int64_t, Int64Type)

        NEBULA_STL_CONVERSION(uint8_t, UInt8Type)

        NEBULA_STL_CONVERSION(uint16_t, UInt16Type)

        NEBULA_STL_CONVERSION(uint32_t, UInt32Type)

        NEBULA_STL_CONVERSION(uint64_t, UInt64Type)

        NEBULA_STL_CONVERSION(float, Fp32Type)

        NEBULA_STL_CONVERSION(double, Fp64Type)

        template<>
        struct ConversionTraits<std::string> : public CTypeTraits<std::string> {
            static turbo::Status append_row(StringBuilder &builder, const std::string &cell) {
                return builder.append(cell);
            }

            static std::string get_entry(const StringArray &array, size_t j) {
                return array.get_string(j);
            }
        };

        /// append cell range elements as a single value to the list builder.
        ///
        /// Cell range will be added to child builder using append_list_values<ValueCType>()
        /// if provided. append_list_values<ValueCType>() has a default implementation, but
        /// it can be specialized by users.
        template<typename ValueCType, typename ListBuilderType, typename Range>
        turbo::Status append_cell_range(ListBuilderType &builder, Range &&cell_range) {
            constexpr bool is_list_builder = std::is_same<ListBuilderType, ListBuilder>::value;
            constexpr bool is_large_list_builder =
                    std::is_same<ListBuilderType, LargeListBuilder>::value;
            static_assert(
                    is_list_builder || is_large_list_builder,
                    "Builder type must be either ListBuilder or LargeListBuilder for appending "
                    "multiple rows.");

            using ChildBuilderType = CBuilderType<ValueCType>;
            TURBO_RETURN_NOT_OK(builder.append());
            auto &value_builder =
                    ::turbo::checked_cast<ChildBuilderType &>(*builder.value_builder());

            // XXX: Remove appended value before returning if status isn't OK?
            return append_list_values<ValueCType>(value_builder, std::forward<Range>(cell_range));
        }

        template<typename ValueCType>
        struct ConversionTraits<std::vector<ValueCType>>
                : public CTypeTraits<std::vector<ValueCType>> {
            static turbo::Status append_row(ListBuilder &builder, const std::vector<ValueCType> &cell) {
                return append_cell_range<ValueCType>(builder, cell);
            }

            static std::vector<ValueCType> get_entry(const ListArray &array, size_t j) {
                using ElementArrayType =
                        typename TypeTraits<typename ConversionTraits<ValueCType>::ArrowType>::ArrayType;

                const ElementArrayType &value_array =
                        ::turbo::checked_cast<const ElementArrayType &>(*array.values());

                std::vector<ValueCType> vec(array.value_length(j));
                for (int64_t i = 0; i < array.value_length(j); i++) {
                    vec[i] =
                            ConversionTraits<ValueCType>::get_entry(value_array, array.value_offset(j) + i);
                }
                return vec;
            }
        };

        template<typename Optional>
        struct ConversionTraits<Optional, enable_if_optional_like<Optional>>
                : public CTypeTraits<typename std::decay<decltype(*std::declval<Optional>())>::type> {
            using OptionalInnerType =
                    typename std::decay<decltype(*std::declval<Optional>())>::type;
            using typename CTypeTraits<OptionalInnerType>::ArrowType;
            using CTypeTraits<OptionalInnerType>::type_singleton;

            static turbo::Status append_row(typename TypeTraits<ArrowType>::BuilderType &builder,
                                           const Optional &cell) {
                if (cell) {
                    return ConversionTraits<OptionalInnerType>::append_row(builder, *cell);
                } else {
                    return builder.append_null();
                }
            }
        };

        /// Build an nebula::Schema based upon the types defined in a std::tuple-like structure.
        ///
        /// While the type information is available at compile-time, we still need to add the
        /// column names at runtime, thus these methods are not constexpr.
        template<typename Tuple, std::size_t N = std::tuple_size<Tuple>::value>
        struct SchemaFromTuple {
            using Element = internal::BareTupleElement<N - 1, Tuple>;

            // Implementations that take a vector-like object for the column names.

            /// Recursively build a vector of nebula::Field from the defined types.
            ///
            /// In most cases MakeSchema is the better entrypoint for the Schema creation.
            static std::vector<std::shared_ptr<Field>> MakeSchemaRecursion(
                    const std::vector<std::string> &names) {
                std::vector<std::shared_ptr<Field>> ret =
                        SchemaFromTuple<Tuple, N - 1>::MakeSchemaRecursion(names);
                auto type = ConversionTraits<Element>::type_singleton();
                ret.push_back(field(names[N - 1], type, internal::is_optional_like<Element>::value));
                return ret;
            }

            /// Build a Schema from the types of the tuple-like structure passed in as template
            /// parameter assign the column names at runtime.
            ///
            /// An example usage of this API can look like the following:
            ///
            /// \code{.cpp}
            /// using TupleType = std::tuple<int, std::vector<std::string>>;
            /// std::shared_ptr<Schema> schema =
            ///   SchemaFromTuple<TupleType>::MakeSchema({"int_column", "list_of_strings_column"});
            /// \endcode
            static std::shared_ptr<Schema> MakeSchema(const std::vector<std::string> &names) {
                return std::make_shared<Schema>(MakeSchemaRecursion(names));
            }

            // Implementations that take a tuple-like object for the column names.

            /// Recursively build a vector of nebula::Field from the defined types.
            ///
            /// In most cases MakeSchema is the better entrypoint for the Schema creation.
            template<typename NamesTuple>
            static std::vector<std::shared_ptr<Field>> MakeSchemaRecursionT(
                    const NamesTuple &names) {
                using std::get;

                std::vector<std::shared_ptr<Field>> ret =
                        SchemaFromTuple<Tuple, N - 1>::MakeSchemaRecursionT(names);
                std::shared_ptr<DataType> type = ConversionTraits<Element>::type_singleton();
                ret.push_back(
                        field(get<N - 1>(names), type, internal::is_optional_like<Element>::value));
                return ret;
            }

            /// Build a Schema from the types of the tuple-like structure passed in as template
            /// parameter assign the column names at runtime.
            ///
            /// An example usage of this API can look like the following:
            ///
            /// \code{.cpp}
            /// using TupleType = std::tuple<int, std::vector<std::string>>;
            /// std::shared_ptr<Schema> schema =
            ///   SchemaFromTuple<TupleType>::MakeSchema({"int_column", "list_of_strings_column"});
            /// \endcode
            template<typename NamesTuple>
            static std::shared_ptr<Schema> MakeSchema(const NamesTuple &names) {
                return std::make_shared<Schema>(MakeSchemaRecursionT<NamesTuple>(names));
            }
        };

        template<typename Tuple>
        struct SchemaFromTuple<Tuple, 0> {
            static std::vector<std::shared_ptr<Field>> MakeSchemaRecursion(
                    const std::vector<std::string> &names) {
                std::vector<std::shared_ptr<Field>> ret;
                ret.reserve(names.size());
                return ret;
            }

            template<typename NamesTuple>
            static std::vector<std::shared_ptr<Field>> MakeSchemaRecursionT(
                    const NamesTuple &names) {
                std::vector<std::shared_ptr<Field>> ret;
                ret.reserve(std::tuple_size<NamesTuple>::value);
                return ret;
            }
        };

        namespace internal {

            template<typename Tuple, std::size_t N = std::tuple_size<Tuple>::value>
            struct CreateBuildersRecursive {
                static turbo::Status create(MemoryPool *pool,
                                          std::vector<std::unique_ptr<ArrayBuilder>> *builders) {
                    using Element = BareTupleElement<N - 1, Tuple>;
                    std::shared_ptr<DataType> type = ConversionTraits<Element>::type_singleton();
                    TURBO_RETURN_NOT_OK(MakeBuilder(pool, type, &builders->at(N - 1)));

                    return CreateBuildersRecursive<Tuple, N - 1>::create(pool, builders);
                }
            };

            template<typename Tuple>
            struct CreateBuildersRecursive<Tuple, 0> {
                static turbo::Status create(MemoryPool *, std::vector<std::unique_ptr<ArrayBuilder>> *) {
                    return turbo::OkStatus();
                }
            };

            template<typename Tuple, std::size_t N = std::tuple_size<Tuple>::value>
            struct RowIterator {
                static turbo::Status append(const std::vector<std::unique_ptr<ArrayBuilder>> &builders,
                                            const Tuple &row) {
                    using std::get;
                    using Element = BareTupleElement<N - 1, Tuple>;
                    using BuilderType =
                            typename TypeTraits<typename ConversionTraits<Element>::ArrowType>::BuilderType;

                    BuilderType &builder =
                            ::turbo::checked_cast<BuilderType &>(*builders[N - 1]);
                    TURBO_RETURN_NOT_OK(ConversionTraits<Element>::append_row(builder, get<N - 1>(row)));

                    return RowIterator<Tuple, N - 1>::append(builders, row);
                }
            };

            template<typename Tuple>
            struct RowIterator<Tuple, 0> {
                static turbo::Status append(const std::vector<std::unique_ptr<ArrayBuilder>> &builders,
                                            const Tuple &row) {
                    return turbo::OkStatus();
                }
            };

            template<typename Tuple, std::size_t N = std::tuple_size<Tuple>::value>
            struct EnsureColumnTypes {
                static turbo::Status Cast(const Table &table, std::shared_ptr<Table> *table_owner,
                                          const compute::CastOptions &cast_options, compute::ExecContext *ctx,
                                          std::reference_wrapper<const ::nebula::Table> *result) {
                    using Element = BareTupleElement<N - 1, Tuple>;
                    std::shared_ptr<DataType> expected_type = ConversionTraits<Element>::type_singleton();

                    if (!table.schema()->field(N - 1)->type()->equals(*expected_type)) {
                        TURBO_MOVE_OR_RAISE(
                                Datum casted,
                                compute::Cast(table.column(N - 1), expected_type, cast_options, ctx));
                        auto new_field = table.schema()->field(N - 1)->with_type(expected_type);
                        TURBO_MOVE_OR_RAISE(*table_owner,
                                            table.set_column(N - 1, new_field, casted.chunked_array()));
                        *result = **table_owner;
                    }

                    return EnsureColumnTypes<Tuple, N - 1>::Cast(result->get(), table_owner, cast_options,
                                                                 ctx, result);
                }
            };

            template<typename Tuple>
            struct EnsureColumnTypes<Tuple, 0> {
                static turbo::Status Cast(const Table &table, std::shared_ptr<Table> *table_owner,
                                          const compute::CastOptions &cast_options, compute::ExecContext *ctx,
                                          std::reference_wrapper<const ::nebula::Table> *result) {
                    return turbo::OkStatus();
                }
            };

            template<typename Range, typename Tuple, std::size_t N = std::tuple_size<Tuple>::value>
            struct TupleSetter {
                static void Fill(const Table &table, Range *rows) {
                    using std::get;
                    using Element = typename std::tuple_element<N - 1, Tuple>::type;
                    using ArrayType =
                            typename TypeTraits<typename ConversionTraits<Element>::ArrowType>::ArrayType;

                    auto iter = rows->begin();
                    const ChunkedArray &chunked_array = *table.column(N - 1);
                    for (int i = 0; i < chunked_array.num_chunks(); i++) {
                        const ArrayType &array =
                                ::turbo::checked_cast<const ArrayType &>(*chunked_array.chunk(i));
                        for (int64_t j = 0; j < array.length(); j++) {
                            get<N - 1>(*iter++) = ConversionTraits<Element>::get_entry(array, j);
                        }
                    }

                    return TupleSetter<Range, Tuple, N - 1>::Fill(table, rows);
                }
            };

            template<typename Range, typename Tuple>
            struct TupleSetter<Range, Tuple, 0> {
                static void Fill(const Table &table, Range *rows) {}
            };

        }  // namespace internal

        template<typename Range>
        turbo::Status TableFromTupleRange(MemoryPool *pool, Range &&rows,
                                          const std::vector<std::string> &names,
                                          std::shared_ptr<Table> *table) {
            using row_type = typename std::iterator_traits<decltype(std::begin(rows))>::value_type;
            constexpr std::size_t n_columns = std::tuple_size<row_type>::value;

            std::shared_ptr<Schema> schema = SchemaFromTuple<row_type>::MakeSchema(names);

            std::vector<std::unique_ptr<ArrayBuilder>> builders(n_columns);
            TURBO_RETURN_NOT_OK(internal::CreateBuildersRecursive<row_type>::create(pool, &builders));

            for (auto const &row: rows) {
                TURBO_RETURN_NOT_OK(internal::RowIterator<row_type>::append(builders, row));
            }

            std::vector<std::shared_ptr<Array>> arrays;
            for (auto const &builder: builders) {
                std::shared_ptr<Array> array;
                TURBO_RETURN_NOT_OK(builder->finish(&array));
                arrays.emplace_back(array);
            }

            *table = Table::create(std::move(schema), std::move(arrays));

            return turbo::OkStatus();
        }

        template<typename Range>
        turbo::Status TupleRangeFromTable(const Table &table, const compute::CastOptions &cast_options,
                                          compute::ExecContext *ctx, Range *rows) {
            using row_type = typename std::decay<decltype(*std::begin(*rows))>::type;
            constexpr std::size_t n_columns = std::tuple_size<row_type>::value;

            if (table.schema()->num_fields() != n_columns) {
                return turbo::invalid_argument_error(
                        "Number of columns in the table does not match the width of the target: ",
                        table.schema()->num_fields(), " != ", n_columns);
            }

            if (std::size(*rows) != static_cast<size_t>(table.num_rows())) {
                return turbo::invalid_argument_error(
                        "Number of rows in the table does not match the size of the target: ",
                        table.num_rows(), " != ", std::size(*rows));
            }

            // Check that all columns have the correct type, otherwise cast them.
            std::shared_ptr<Table> table_owner;
            std::reference_wrapper<const ::nebula::Table> current_table(table);

            TURBO_RETURN_NOT_OK(internal::EnsureColumnTypes<row_type>::Cast(
                    table, &table_owner, cast_options, ctx, &current_table));

            internal::TupleSetter<Range, row_type>::Fill(current_table.get(), rows);

            return turbo::OkStatus();
        }

    }  // namespace stl
}  // namespace nebula
